import numpy as np
import torch
from torch import nn
import colorsys
import time

# 导入test_strip.py中的函数
from test_strip import create_display_command, send_command

# 你的LED参数
GROUP_ADDRESS = 0
DEVICE_ADDRESS = 1
PORT = 0
LED_TYPE = 0x01  # WS2812B
LED_COUNT = 30   # 你实际的灯珠数量
REPEAT = 1

class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.input_hidden = nn.Linear(input_size, hidden_size)
        self.hidden_output = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        hidden = torch.sigmoid(self.input_hidden(x))
        output = self.hidden_output(hidden)
        return output

def predict(x, model):
    x_torch = torch.tensor(x, dtype=torch.float32).to(device)
    output_activations = model(x_torch)
    return output_activations.detach().cpu().numpy()

def main():
    arousal = np.random.uniform(-1, 1)
    valence = np.random.uniform(-1, 1)
    change_range = 0.1

    while True:
        delta_arousal = np.random.uniform(-change_range, change_range)
        delta_valence = np.random.uniform(-change_range, change_range)
        arousal = np.clip(arousal + delta_arousal, -1, 1)
        valence = np.clip(valence + delta_valence, -1, 1)
        X_test = np.array([[arousal, valence]])
        H_predictions = predict(X_test, model_H).flatten()
        S_predictions = predict(X_test, model_S).flatten()
        
        B_constant = 1
        print(f"Hue: {H_predictions[0]}, Saturation: {S_predictions[0]}")
        # color = np.clip(colorsys.hsv_to_rgb(H_predictions[0], S_predictions[0], B_constant), 0, 1)
        color=colorsys.rgb_to_hsv(H_predictions[0], S_predictions[0], B_constant)
        r, g, b = int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)

        print(f"Arousal: {arousal:.2f}, Valence: {valence:.2f}, RGB: ({r}, {g}, {b})")

        # 生成指令并发送到灯带
        colors = [(r, g, b)] * LED_COUNT
        cmd = create_display_command(
            GROUP_ADDRESS,
            DEVICE_ADDRESS,
            PORT,
            LED_TYPE,
            REPEAT,
            colors
        )
        send_command(cmd)

        time.sleep(0.5)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_H = SimpleNN(2, 7, 1)  # 按你实际的参数
model_H.load_state_dict(torch.load('trained of H_model.pth', map_location=device))
model_H.to(device)
model_H.eval()

model_S = SimpleNN(2, 7, 1)  # 按你实际的参数
model_S.load_state_dict(torch.load('trained of S_model.pth', map_location=device))
model_S.to(device)
model_S.eval()

if __name__ == '__main__':
    main()